import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import os

# 定义与训练时相同的模型结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

def load_model(model_path):
    # 检查CUDA是否可用
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # 加载模型
    model = Net().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    
    return model, device

def preprocess_image(image_path):
    # 打开图像
    image = Image.open(image_path).convert('L')  # 转换为灰度图
    
    # 转换为MNIST格式 (28x28)
    image = image.resize((28, 28))
    
    # 转换为张量并规范化
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    # 增加批次维度
    image_tensor = transform(image).unsqueeze(0)
    
    return image_tensor, image

def predict_digit(model, image_tensor, device):
    # 将图像张量移至适当的设备
    image_tensor = image_tensor.to(device)
    
    # 预测
    with torch.no_grad():
        output = model(image_tensor)
        pred = output.argmax(dim=1, keepdim=True)
        probs = F.softmax(output, dim=1)
    
    return pred.item(), probs[0].cpu().numpy()

def display_prediction(image, prediction, probabilities):
    # 显示图像和预测结果
    plt.figure(figsize=(10, 6))
    
    # 显示原始图像
    plt.subplot(1, 2, 1)
    plt.imshow(image, cmap='gray')
    plt.title(f"预测结果: {prediction}")
    plt.axis('off')
    
    # 显示概率分布
    plt.subplot(1, 2, 2)
    plt.bar(range(10), probabilities)
    plt.xticks(range(10))
    plt.title("各数字的概率")
    plt.xlabel("数字")
    plt.ylabel("概率")
    
    plt.tight_layout()
    plt.savefig('prediction_result.png')
    print("预测结果已保存为 'prediction_result.png'")
    plt.show()

def main():
    model_path = "mnist_cnn.pt"
    
    if not os.path.exists(model_path):
        print(f"错误: 找不到模型文件 '{model_path}'")
        print("请先运行 train.py 训练模型")
        return
    
    # 加载模型
    model, device = load_model(model_path)
    
    # 用户输入图像路径
    while True:
        image_path = input("请输入手写数字图像路径(输入'exit'退出): ")
        
        if image_path.lower() == 'exit':
            break
            
        if not os.path.exists(image_path):
            print(f"错误: 找不到图像文件 '{image_path}'")
            continue
        
        try:
            # 预处理图像
            image_tensor, image = preprocess_image(image_path)
            
            # 预测数字
            prediction, probabilities = predict_digit(model, image_tensor, device)
            
            # 显示预测结果
            print(f"预测的数字是: {prediction}")
            print(f"预测概率: {probabilities[prediction]:.4f}")
            
            # 显示图像和预测结果
            display_prediction(image, prediction, probabilities)
            
        except Exception as e:
            print(f"处理图像时出错: {e}")
    
    print("预测完成!")

if __name__ == "__main__":
    main() 